'''Remote machine control via SSH protocol and public-private key encryption.'''

import pexpect

def command(command, password, 
        username='usf_ubcslice5',
        keyfile='id_rsa', 
        hostname='planetlab4.cs.uiuc.edu',
        ):
    '''Executes a command on a remote machine via ssh.'''
    try:
        # -t spawns a pseudo-tty, this fixes some weird ssh-deploy
        # only glitch I ran into
        ssh = pexpect.spawn('ssh -i %s -t %s@%s "%s"' % 
            (keyfile, username, hostname, command))


        SWITCH = [
            "Are you sure you want to continue connecting",
            "id_rsa': ",
            pexpect.EOF
        ]

        i=ssh.expect(SWITCH)
        if i==0:
            print "ssh.command(): hostname %s added to ~/.ssh/known_hosts!" % (hostname)
            ssh.sendline('yes')
            i=ssh.expect(SWITCH)

        if i==1:
            ssh.sendline(password)
            ssh.expect(pexpect.EOF, timeout=60*10)
        elif i==2:
            print "ssh.command(): command failed to execute due to timeout."
            print "Returned:",ssh.before
            return None

        return ssh.before

    except:
        print 'Error:',ssh.before
        return None


def upload(src_file, dst_file, password,
        username='usf_ubcslice5',
        keyfile='/home/kevin/p01/id_rsa', 
        hostname='planetlab4.cs.uiuc.edu',
        ):
    '''Copies a file to a remote machine via scp.'''
    try:
        ssh = pexpect.spawn('scp -i %s %s %s@%s:%s' % 
            (keyfile, src_file, username, hostname, dst_file))


        SWITCH = [
            "Are you sure you want to continue connecting",
            "id_rsa': ",
            pexpect.EOF
        ]

        i=ssh.expect(SWITCH)
        if i==0:
            print "ssh.upload() hostname %s added to ~/.ssh/known_hosts!" % (hostname)
            ssh.sendline('yes')
            i=ssh.expect(SWITCH)
        if i==1:
            ssh.sendline(password)
            ssh.expect(pexpect.EOF, timeout=60*30) # 30 minute timeout
        elif i==2:
            print "ssh.upload() failed due to timeout."
            return None

        return ssh.before
    except: return None

